#!/usr/bin/env python3

import pickle

#  定义数据包类型
#  客户端连接服务器请求包
CONNECT = 1
#  错误数据包
ERROR = 2
#  连接成功包
CONNECT_ACK = 3
#  数据包
DATA = 4
#  数据回应包
DATA_ACK = 5
#  断开连接包
SHUTDOWN = 6

MAXSIZE = 1024

#  错误信息字典
errorMessage = {
    0: '用户名不存在',
    1: '密码验证失败',
    2: '文件不存在或者文件异常',
    3: '非法数据包'
}

def get(data, types=0):
    d = pickle.loads(data)
    if (type(d) != dict
            or 'type' not in d
            or (types != 0 and d['type'] != types)):
        return None
    return d

def put(data, types):
    if data == None:
        d = {}
    elif type(data) != dict:
        d = {}
        d['data'] = data
    else:
        d = data

    d['type'] = types

    return pickle.dumps(d)

def send(sd, data, addr, types):
    sendData = put(data, types)
    l = len(sendData)
    return l if sd.sendto(sendData, addr) == l else None

def recv(sd, types=0, timeout=None):
    sd.settimeout(timeout)
    try:
        data, addr = sd.recvfrom(MAXSIZE)
    except Exception as e:
        return None
    finally:
        sd.settimeout(None)

    r = get(data, types)
    return (r, addr) if r else None

